package library

type Trie struct {
	root *TrieNode
}

var AcCtx = NewTrie()

type TrieNode struct {
	End      bool
	Next     map[rune]*TrieNode
	FailNode *TrieNode
	Depth    int
}
type KeyWord struct {
	KeyWord string `json:"keyword"`
	Start   int64  `json:"start"`
	End     int64  `json:"end"`
}

func NewTrie() Trie {
	var r Trie
	r.root = r.newTrieNode(0)
	return r
}
func init() {
	buildAc(&AcCtx)
}
func buildAc(tire *Trie) {
	keywords := []string{"麻辣", "鲜香一品", "诸葛", "万州的", "州烤鱼"}
	for _, v := range keywords {
		tire.Add(v)
	}
	tire.BuildFaild()
}

func (t *Trie) newTrieNode(depth int) *TrieNode {
	n := new(TrieNode)
	n.Next = make(map[rune]*TrieNode)
	n.Depth = depth
	return n
}

func (t *Trie) Add(keyword string) {
	if keyword == "" {
		return
	}
	node := t.root
	i := 1
	for _, v := range keyword {
		if _, exists := node.Next[v]; !exists {
			node.Next[v] = t.newTrieNode(i)
		}
		node = node.Next[v]
		i++
	}
	node.End = true
}

func (t *Trie) Query(content []rune) (results []KeyWord) {
	iter := t.root
	var start, end int
	for i, c := range content {
		_, ok := iter.Next[c]
		if !ok && iter != t.root {
			iter = iter.FailNode
		}
		if _, ok := iter.Next[c]; !ok {
			iter = t.root
			continue
		}
		start = i - iter.Depth
		iter = iter.Next[c]
		if iter.End {
			end = i
			KeyWord := KeyWord{
				KeyWord: string(content[start : end+1]),
				Start:   int64(start),
				End:     int64(end),
			}
			results = append(results, KeyWord)
		}
	}
	return
}
func (t *Trie) KeywordsCount(content []rune) (results map[string]int) {
	iter := t.root
	results = make(map[string]int)
	var start, end int
	for i, c := range content {
		_, ok := iter.Next[c]
		if !ok && iter != t.root {
			iter = iter.FailNode
		}
		if _, ok := iter.Next[c]; !ok {
			iter = t.root
			continue
		}
		start = i - iter.Depth
		iter = iter.Next[c]
		if iter.End {
			end = i
			keyword := string(content[start : end+1])
			if _, ok = results[keyword]; ok {
				results[keyword]++
			} else {
				results[keyword] = 1
			}
		}
	}
	return
}

func (t *Trie) BuildFaild() {
	queue := []*TrieNode{t.root}
	for len(queue) > 0 {
		parent := queue[0]
		queue = queue[1:]
		for char, child := range parent.Next {
			if parent == t.root {
				child.FailNode = t.root
			} else {
				failAcNode := parent.FailNode
				for failAcNode != nil {
					if _, ok := failAcNode.Next[char]; ok {
						child.FailNode = failAcNode.Next[char]
						break
					}
					failAcNode = failAcNode.FailNode
				}
				if failAcNode == nil {
					child.FailNode = t.root
				}
			}
			queue = append(queue, child)
		}

	}
}
